{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 手写数字训练集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torch.utils import data\n",
    "from torch import nn\n",
    "import torch\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 512\n",
    "NUM_EPOCHS = 10\n",
    "mnist_train = torchvision.datasets.MNIST(root='./data',transform=transforms.ToTensor(),train=True,download=True)\n",
    "mnist_test = torchvision.datasets.MNIST(root='./data',transform=transforms.ToTensor(),train=False,download=True)\n",
    "mnist_train = data.DataLoader(mnist_train,batch_size=BATCH_SIZE)\n",
    "mnist_test = data.DataLoader(mnist_test,batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = nn.Sequential(\n",
    "    nn.Flatten(),\n",
    "    nn.Linear(28 * 28,10)\n",
    ")\n",
    "\n",
    "def init_weight(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        nn.init.normal_(m.weight,0.01)\n",
    "\n",
    "net.apply(init_weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#对于分类任务，损失函数用CrossEntropyLoss\n",
    "loss = nn.CrossEntropyLoss(reduction='none')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#优化器\n",
    "trainer = torch.optim.SGD(net.parameters(),lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def accuracy(y_hat,y):\n",
    "    y_hat = y_hat.argmax(axis=1)\n",
    "    return float((y_hat.type(y.dtype) == y).type(y.dtype).sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(NUM_EPOCHS):\n",
    "    for x,y in mnist_train:\n",
    "        y_hat = net(x)\n",
    "        l = loss(y_hat,y)\n",
    "        trainer.zero_grad()\n",
    "        l.mean().backward()\n",
    "        trainer.step()\n",
    "    acc, losses, num = 0.0,0.0,0.0\n",
    "    with torch.no_grad():\n",
    "        for X,y in mnist_train:\n",
    "            acc += accuracy(net(X),y)\n",
    "            losses += float(loss(net(X),y).sum())\n",
    "            num += y.numel()\n",
    "    print(\"train loss\",losses / num,\"train acc\",acc /num)\n",
    "    acc, losses, num = 0.0,0.0,0.0\n",
    "    with torch.no_grad():\n",
    "        for X,y in mnist_test:\n",
    "            acc += accuracy(net(X),y)\n",
    "            losses += float(loss(net(X),y).sum())\n",
    "            num += y.numel()\n",
    "    print(\"test loss\",losses / num,\"test acc\",acc /num)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c = next(iter(mnist_test))[0][0]\n",
    "next(iter(mnist_test))[1][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net(c)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d0l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
